package query

import (
	"bufio"
	"fmt"
	"regexp"
	"strconv"
	"strings"
)

// InspurQuery Inspur GPU查询实现
type InspurQuery struct{}

func init() {
	RegisterQueryer("inspur", func() GPUQuery { return &InspurQuery{} })
}

func (q *InspurQuery) Query() ([]GPUInfo, error) {
	// 获取设备列表
	devices, err := q.getDevices()
	if err != nil {
		return nil, fmt.Errorf("获取设备列表失败: %v", err)
	}

	var gpus []GPUInfo
	for _, device := range devices {
		// 获取设备使用率和内存信息
		utilization, memoryUsed, memoryTotal, err := q.getDeviceInfo(device)
		if err != nil {
			return nil, fmt.Errorf("获取设备%d信息失败: %v", device, err)
		}

		gpus = append(gpus, GPUInfo{
			Index:       device,
			Utilization: utilization,
			MemoryUsed:  memoryUsed,
			MemoryTotal: memoryTotal,
		})
	}

	return gpus, nil
}

// getDevices 获取设备列表
func (q *InspurQuery) getDevices() ([]int, error) {
	// 使用查询命令获取设备信息
	output, err := runCommand("hxsmi", "query", "-s", "MEMORY,UTILIZATION")
	if err != nil {
		return nil, err
	}

	var devices []int
	scanner := bufio.NewScanner(strings.NewReader(output))
	cardRegex := regexp.MustCompile(`Card\s+(\d+)`)

	for scanner.Scan() {
		line := scanner.Text()
		if matches := cardRegex.FindStringSubmatch(line); matches != nil {
			if deviceID, err := strconv.Atoi(matches[1]); err == nil {
				devices = append(devices, deviceID)
			}
		}
	}

	if len(devices) == 0 {
		return nil, fmt.Errorf("未找到有效的GPU设备")
	}

	return devices, nil
}

// getDeviceInfo 获取设备使用率和内存信息
func (q *InspurQuery) getDeviceInfo(device int) (int, int, int, error) {
	// 获取完整信息
	output, err := runCommand("hxsmi", "query", "-s", "MEMORY,UTILIZATION", "-i", strconv.Itoa(device))
	if err != nil {
		return 0, 0, 0, err
	}

	// 解析使用率信息
	utilization, err := q.parseUtilization(output)
	if err != nil {
		return 0, 0, 0, err
	}

	// 解析内存信息
	memoryUsed, memoryTotal, err := q.parseMemoryInfo(output)
	if err != nil {
		return 0, 0, 0, err
	}

	return utilization, memoryUsed, memoryTotal, nil
}

// parseUtilization 解析使用率信息
func (q *InspurQuery) parseUtilization(output string) (int, error) {
	scanner := bufio.NewScanner(strings.NewReader(output))
	utilRegex := regexp.MustCompile(`GPU\s+:\s+(\d+) %`)

	for scanner.Scan() {
		line := scanner.Text()
		if matches := utilRegex.FindStringSubmatch(line); matches != nil {
			if utilization, err := strconv.Atoi(matches[1]); err == nil {
				return utilization, nil
			}
		}
	}

	return 0, fmt.Errorf("未找到GPU使用率信息")
}

// parseMemoryInfo 解析内存信息
func (q *InspurQuery) parseMemoryInfo(output string) (int, int, error) {
	scanner := bufio.NewScanner(strings.NewReader(output))
	memoryUsedRegex := regexp.MustCompile(`Used\s+:\s+(\d+) MiB`)
	memoryTotalRegex := regexp.MustCompile(`Total\s+:\s+(\d+) MiB`)

	var used, total int
	foundUsed, foundTotal := false, false

	for scanner.Scan() {
		line := scanner.Text()
		if matches := memoryUsedRegex.FindStringSubmatch(line); matches != nil && !foundUsed {
			if value, err := strconv.Atoi(matches[1]); err == nil {
				used = value
				foundUsed = true
			}
		}
		if matches := memoryTotalRegex.FindStringSubmatch(line); matches != nil && !foundTotal {
			if value, err := strconv.Atoi(matches[1]); err == nil {
				total = value
				foundTotal = true
			}
		}
		if foundUsed && foundTotal {
			break
		}
	}

	if !foundUsed || !foundTotal {
		return 0, 0, fmt.Errorf("未找到完整的内存信息")
	}

	return used, total, nil
}
